{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -e ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update this for your data path.\n",
    "instance_data_dir = \"/Volumes/ml/datasets/test_datasets/single_image_dataset\"\n",
    "pretrained_model_name_or_path = \"black-forest-labs/FLUX.1-dev\"\n",
    "# Your public model name after it's pushed to the hub.\n",
    "hub_model_id = \"simpletuner-lora\"\n",
    "tracker_project_name = \"flux-training\"\n",
    "\n",
    "# Validation prompt\n",
    "validation_prompt = \"A photo-realistic image of a cat\"\n",
    "\n",
    "train_batch_size = 1\n",
    "learning_rate = 1e-4\n",
    "\n",
    "# choices: int8-quanto, fp8-quanto, no_change (mac and a100/h100 users get int4 and int2 as well)\n",
    "base_model_precision = \"no_change\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lycoris_config = {\n",
    "    \"algo\": \"lokr\",\n",
    "    \"multiplier\": 1.0,\n",
    "    \"linear_dim\": 10000,\n",
    "    \"linear_alpha\": 1,\n",
    "    \"factor\": 12,\n",
    "    \"apply_preset\": {\n",
    "        \"target_module\": [\n",
    "            \"Attention\",\n",
    "            \"FeedForward\"\n",
    "        ],\n",
    "        \"module_algo_map\": {\n",
    "            \"Attention\": {\n",
    "                \"factor\": 12\n",
    "            },\n",
    "            \"FeedForward\": {\n",
    "                \"factor\": 6\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "# write to config/lycoris_config.json\n",
    "import json\n",
    "with open(\"config/lycoris_config.json\", \"w\") as f:\n",
    "    json.dump(lycoris_config, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_config = {\n",
    "    \"mixed_precision\":\"bf16\",\n",
    "    \"model_type\":\"lora\",\n",
    "    \"pretrained_model_name_or_path\":pretrained_model_name_or_path,\n",
    "    \"gradient_checkpointing\":True,\n",
    "    \"cache_dir\": \"cache\",\n",
    "    \"set_grads_to_none\":True,\n",
    "    \"gradient_accumulation_steps\":1,\n",
    "    \"resume_from_checkpoint\":\"latest\",\n",
    "    \"snr_gamma\":5,\n",
    "    \"num_train_epochs\":0,\n",
    "    \"max_train_steps\":10000,\n",
    "    \"metadata_update_interval\":65,\n",
    "    \"optimizer\":\"adamw_bf16\",\n",
    "    \"learning_rate\":learning_rate,\n",
    "    \"lr_scheduler\":\"polynomial\",\n",
    "    \"seed\":42,\n",
    "    \"lr_warmup_steps\":100,\n",
    "    \"output_dir\":\"output/models\",\n",
    "    \"aspect_bucket_rounding\":2,\n",
    "    \"inference_scheduler_timestep_spacing\":\"trailing\",\n",
    "    \"training_scheduler_timestep_spacing\":\"trailing\",\n",
    "    \"report_to\":\"wandb\",\n",
    "    \"lr_end\":1e-8,\n",
    "    \"compress_disk_cache\":True,\n",
    "    \"push_to_hub\":True,\n",
    "    \"hub_model_id\":hub_model_id,\n",
    "    \"push_checkpoints_to_hub\":True,\n",
    "    \"model_family\":\"flux\",\n",
    "    \"disable_benchmark\":False,\n",
    "    \"train_batch\":train_batch_size,\n",
    "    \"max_workers\":32,\n",
    "    \"read_batch_size\":25,\n",
    "    \"write_batch_size\":64,\n",
    "    \"caption_dropout_probability\":0.1,\n",
    "    \"torch_num_threads\":8,\n",
    "    \"image_processing_batch_size\":32,\n",
    "    \"vae_batch_size\":4,\n",
    "    \"validation_prompt\":validation_prompt,\n",
    "    \"num_validation_images\":1,\n",
    "    \"validation_num_inference_steps\":20,\n",
    "    \"validation_seed\":42,\n",
    "    \"minimum_image_size\":0,\n",
    "    \"resolution\":1024,\n",
    "    \"validation_resolution\":\"1024x1024\",\n",
    "    \"resolution_type\":\"pixel_area\",\n",
    "    \"lycoris_config\":\"config/lycoris_config.json\",\n",
    "    \"lora_type\":\"lycoris\",\n",
    "    \"base_model_precision\":base_model_precision,\n",
    "    \"checkpointing_steps\":500,\n",
    "    \"checkpoints_total_limit\":5,\n",
    "    \"validation_steps\":500,\n",
    "    \"tracker_run_name\":hub_model_id,\n",
    "    \"tracker_project_name\":tracker_project_name,\n",
    "    \"validation_guidance\":3.0,\n",
    "    \"validation_guidance_real\":1.0,\n",
    "    \"validation_guidance_rescale\":0.0,\n",
    "    \"validation_negative_prompt\":\"blurry, cropped, ugly\",\n",
    "}\n",
    "# write to config/config.json\n",
    "with open(\"config/config.json\", \"w\") as f:\n",
    "    json.dump(training_config, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader_config = [\n",
    "    {\n",
    "        \"id\": \"my-dataset-512\",\n",
    "        \"type\": \"local\",\n",
    "        \"instance_data_dir\": instance_data_dir,\n",
    "        \"crop\": False,\n",
    "        \"crop_style\": \"random\",\n",
    "        \"minimum_image_size\": 128,\n",
    "        \"resolution\": 512,\n",
    "        \"resolution_type\": \"pixel_area\",\n",
    "        \"repeats\": \"4\",\n",
    "        \"metadata_backend\": \"discovery\",\n",
    "        \"caption_strategy\": \"filename\",\n",
    "        \"cache_dir_vae\": \"cache/vae-512\"\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"my-dataset-1024\",\n",
    "        \"type\": \"local\",\n",
    "        \"instance_data_dir\": instance_data_dir,\n",
    "        \"crop\": False,\n",
    "        \"crop_style\": \"random\",\n",
    "        \"minimum_image_size\": 128,\n",
    "        \"resolution\": 1024,\n",
    "        \"resolution_type\": \"pixel_area\",\n",
    "        \"repeats\": \"4\",\n",
    "        \"metadata_backend\": \"discovery\",\n",
    "        \"caption_strategy\": \"filename\",\n",
    "        \"cache_dir_vae\": \"cache/vae-1024\"\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"my-dataset-512-crop\",\n",
    "        \"type\": \"local\",\n",
    "        \"instance_data_dir\": instance_data_dir,\n",
    "        \"crop\": False,\n",
    "        \"crop_style\": \"random\",\n",
    "        \"minimum_image_size\": 128,\n",
    "        \"resolution\": 512,\n",
    "        \"resolution_type\": \"pixel_area\",\n",
    "        \"repeats\": \"4\",\n",
    "        \"metadata_backend\": \"discovery\",\n",
    "        \"caption_strategy\": \"filename\",\n",
    "        \"cache_dir_vae\": \"cache/vae-512-crop\"\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"my-dataset-1024-crop\",\n",
    "        \"type\": \"local\",\n",
    "        \"instance_data_dir\": instance_data_dir,\n",
    "        \"crop\": False,\n",
    "        \"crop_style\": \"random\",\n",
    "        \"minimum_image_size\": 128,\n",
    "        \"resolution\": 1024,\n",
    "        \"resolution_type\": \"pixel_area\",\n",
    "        \"repeats\": \"4\",\n",
    "        \"metadata_backend\": \"discovery\",\n",
    "        \"caption_strategy\": \"filename\",\n",
    "        \"cache_dir_vae\": \"cache/vae-1024-crop\"\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"text-embed-cache\",\n",
    "        \"dataset_type\": \"text_embeds\",\n",
    "        \"default\": True,\n",
    "        \"type\": \"local\",\n",
    "        \"cache_dir\": \"cache/text\"\n",
    "    }\n",
    "]\n",
    "# write to config/multidatabackend.json\n",
    "import json\n",
    "with open(\"config/multidatabackend.json\", \"w\") as f:\n",
    "    json.dump(dataloader_config, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import logging.config\n",
    "from os import environ\n",
    "\n",
    "logging.getLogger(\"DeepSpeed\").setLevel(\"ERROR\")\n",
    "logging.getLogger(\"torch.distributed.elastic.multiprocessing.redirects\").setLevel(\"ERROR\")\n",
    "\n",
    "logging.config.dictConfig({\n",
    "    \"version\": 1,\n",
    "    \"disable_existing_loggers\": True,\n",
    "})\n",
    "\n",
    "environ[\"ACCELERATE_LOG_LEVEL\"] = \"WARNING\"\n",
    "\n",
    "from simpletuner.helpers import log_format\n",
    "from simpletuner.helpers.training.multi_process import _get_rank\n",
    "from simpletuner.helpers.training.state_tracker import StateTracker\n",
    "from simpletuner.helpers.training.trainer import Trainer\n",
    "\n",
    "logger = logging.getLogger(\"SimpleTuner\")\n",
    "logger.setLevel(environ.get(\"SIMPLETUNER_LOG_LEVEL\", \"INFO\") if _get_rank() == 0 else \"ERROR\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from simpletuner.helpers.configuration.json_file import normalize_args\n",
    "import os\n",
    "os.environ['CONFIG_BACKEND'] = 'cmd'\n",
    "os.environ['ENV'] = 'default'\n",
    "StateTracker.set_config_path('config/')\n",
    "loaded_config = normalize_args(training_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing\n",
    "\n",
    "try:\n",
    "    multiprocessing.set_start_method('fork')\n",
    "except Exception as e:\n",
    "    logger.error(\n",
    "        \"Failed to set the multiprocessing start method to 'fork'. Unexpected behaviour such as high memory overhead or poor performance may result.\"\n",
    "        f\"\\nError: {e}\"\n",
    "    )\n",
    "\n",
    "try:\n",
    "    trainer = Trainer(\n",
    "        loaded_config,\n",
    "        exit_on_error=True,\n",
    "    )\n",
    "except Exception as e:\n",
    "    import traceback\n",
    "    logger.error(f\"Failed to create Trainer: {e}, {traceback.format_exc()}\")\n",
    "    raise e\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    trainer.configure_webhook()\n",
    "    trainer.init_noise_schedule()\n",
    "    trainer.init_seed()\n",
    "\n",
    "    trainer.init_huggingface_hub()\n",
    "except Exception as e:\n",
    "    logger.error(f\"Failed to configure Trainer: {e}\")\n",
    "    raise e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    trainer.init_preprocessing_models()\n",
    "    trainer.init_precision(preprocessing_models_only=True)\n",
    "except Exception as e:\n",
    "    logger.error(f\"Failed to initialize preprocessing models: {e}\")\n",
    "    raise e\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    trainer.init_data_backend()\n",
    "except Exception as e:\n",
    "    logger.error(f\"Failed to initialize data backend: {e}\")\n",
    "    raise e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_unload_text_encoder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_unload_vae()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_load_base_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_controlnet_model()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_tread_model()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_precision()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_freeze_models()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_trainable_peft_adapter()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_ema_model()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_precision(ema_only=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.move_models(destination=\"accelerator\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_distillation()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_validations()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.enable_sageattention_inference()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_benchmark_base_model()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.disable_sageattention_inference()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.resume_and_prepare()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.init_trackers()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
